'''
Created on Feb 13, 2011

@author: oabalbin
'''


import os
import sys
import subprocess

from optparse import OptionParser
from exome.jobs.base import JOB_SUCCESS, JOB_ERROR
from exome.jobs.job_runner import qsub_cac, qsub_loc, run_local
from exome.jobs.config import ExomePipelineConfig, ExomeAnalysisConfig

# Global Variables
NODE_MEM=45000.0
NODE_CORES=12
SINGLE_CORE=1
MEM_PER_CORE= int(float(NODE_MEM) / NODE_CORES)
# wt=walltime
WT_SHORT= "10:00:00" #24.00.00
WT_LONG= "60:00:00" #"100:00:00"
WT_ELONG="100:00:00"


def tabix_vcf_files(vcf_file):
    '''
    It first compress the vcf file
    And then it creates a tabix index of the vcf file
    '''
    bgzip=['bgzip',vcf_file]
    bgzip=",".join(bgzip).replace(',',' ')
    tabix=['tabix','-p', 'vcf', vcf_file.replace('.vcf','.vcf.gz')]
    tabix=",".join(tabix).replace(',',' ')
    
    return bgzip+'\n'+tabix

def vcf_isec(A_vcf_file, output_vcf_file, list_vcf_files, path_to_vcftools, do_complement=False):
    '''
    vcf-isec [OPTIONS] file1.vcf file2.vcf
    -c : do the complent. default None
    -f : foce the intersection evetough vcf 
        files with different columns
    '''
    if do_complement:
        complement=['-c','-f']
    else:
        complement=['-o','-f']
        
    vcftools_command=path_to_vcftools+'bin/vcf-isec'
    args=[vcftools_command]+\
          complement+\
          [A_vcf_file]+\
          list_vcf_files+\
          ['>', output_vcf_file]
    args=",".join(args).replace(',',' ')    
    return args


def variant_isec(analysis, configrun, query_vcf_file, jobrunfunc, do_complement, output_vcf_file=None):

    # Parameters
    extra_mem, num_cores = configrun.gatk_use_mem, configrun.gatk_num_cores
    path_to_vcftools, path_to_picard, path_to_sam = configrun.vcftools_path, configrun.picard_path, \
                                                configrun.samtools_path
    genomes = configrun.genomes['human']
    ref_genome, snpdb_vcf = genomes.gatk_ref_genome, genomes.snpdb
    hapmap_vcf, tgk_vcf = genomes.hapmap, genomes.OneKgenomes 
    my_email=configrun.email_addresses
    
    ###########    
    jobn='knownSNPs_'
    jobidvcf=None
    if output_vcf_file is None:
        output_vcf_file=query_vcf_file.replace('.vcf','.knownSNPs.isec.vcf')
    #do_complement=True
    list_vcf_files=[hapmap_vcf, tgk_vcf, snpdb_vcf]
    ###########

    # check if file is already in gz file, also check for the tabix file.
    
    tabix=[]
    
    # check if query_vfc file is compressed and tabix index
    if os.path.isfile(query_vcf_file.replace('.vcf','.vcf.gz')) and \
       os.path.isfile(query_vcf_file.replace('.vcf','.vcf.gz.tbi')):
        tabix.append('echo vcf File %s ready for intersection'%query_vcf_file)
        query_vcf_file_gz = query_vcf_file.replace('.vcf','.vcf.gz')
    else:
        tabix.append(tabix_vcf_files(query_vcf_file))
        query_vcf_file_gz = query_vcf_file.replace('.vcf','.vcf.gz')

    # Check if all other files are compressed and tabix index
    ready_vcf_files=[]
    for vcf_file_name in list_vcf_files:
        if os.path.isfile(vcf_file_name.replace('.vcf','.vcf.gz')) and \
           os.path.isfile(vcf_file_name.replace('.vcf','.vcf.gz.tbi')):
            tabix.append('echo vcf File %s ready for intersection'%vcf_file_name)
            ready_vcf_files.append(vcf_file_name.replace('.vcf','.vcf.gz'))
        else:
            tabix.append(tabix_vcf_files(vcf_file_name))
            ready_vcf_files.append(vcf_file_name.replace('.vcf','.vcf.gz'))

    
    # Intersect vfc file with known snps.
    tabix.append(vcf_isec(query_vcf_file_gz, output_vcf_file, ready_vcf_files, path_to_vcftools, do_complement))
    
    cmd=",".join(tabix).replace(',','\n')
    jobidvcf = jobrunfunc(jobn+'isec', cmd, SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=None, 
                           deps=jobidvcf, stdout=None, email_addresses=my_email)
    
    return jobidvcf


def knownSNPs_isec(analysis, configrun, vcf_file_list, jobrunfunc, do_complement):
    '''
    '''
    for vcf_file_name in vcf_file_list:
        variant_isec(analysis, configrun, vcf_file_name, jobrunfunc,do_complement)
        
def SNPs_isec(query_vcf_file, others_vcf_files, configrun, jobrunfunc, do_complement, output_vcf_file=None, depends=None):
    '''
    '''
    # Parameters
    extra_mem, num_cores = configrun.gatk_use_mem, configrun.gatk_num_cores
    path_to_vcftools, path_to_picard, path_to_sam = configrun.vcftools_path, configrun.picard_path, \
                                                configrun.samtools_path
    my_email=configrun.email_addresses
    
    ###########    
    jobn='SNPs_'
    
    if depends is not None:
        prev_deps=depends[:]
    else:
        prev_deps=[]

    
    if output_vcf_file is None:
        output_vcf_file=query_vcf_file.replace('.vcf','.SNPs.isec.vcf')
    ###########
 
    # check if query_vfc file is compressed and tabix index
    tabix=[]
    if os.path.isfile(query_vcf_file.replace('.vcf','.vcf.gz')) and \
    os.path.isfile(query_vcf_file.replace('.vcf','.vcf.gz.tbi')):
        tabix.append('echo vcf File %s ready for intersection'%query_vcf_file)
        query_vcf_file_gz=query_vcf_file.replace('.vcf','.vcf.gz')
    else:
        tabix.append(tabix_vcf_files(query_vcf_file))
        query_vcf_file_gz=query_vcf_file.replace('.vcf','.vcf.gz')
    
    # Check if all other files are compressed and tabix index
    ready_vcf_files=[]
    for vcf_file_name in others_vcf_files:
        if os.path.isfile(vcf_file_name.replace('.vcf','.vcf.gz')) and \
        os.path.isfile(vcf_file_name.replace('.vcf','.vcf.gz.tbi')):
            tabix.append('echo vcf File %s ready for intersection'%vcf_file_name)
            ready_vcf_files.append(vcf_file_name.replace('.vcf','.vcf.gz'))
        else:
            tabix.append(tabix_vcf_files(vcf_file_name))
            ready_vcf_files.append(vcf_file_name.replace('.vcf','.vcf.gz'))
    
    # Intersect vfc file with known snps.
    tabix.append( vcf_isec(query_vcf_file_gz, output_vcf_file, ready_vcf_files, 
                           path_to_vcftools, do_complement))
    
    cmd=",".join(tabix).replace(',','\n')
    jobidvcf = jobrunfunc(jobn+'isec', cmd, SINGLE_CORE, cwd=None, walltime=WT_SHORT, pmem=None, 
                           deps=prev_deps, stdout=None, email_addresses=my_email)
    
    return jobidvcf

    


if __name__ == '__main__':
        
    optionparser = OptionParser("usage: %prog [options] ")
    
    optionparser.add_option("-r", "--config_file", dest="config_file",
                            help="file with run configuration")
    optionparser.add_option("-a", "--analysis_file", dest="analysis_file",
                            help="file with experiment configuration")
    optionparser.add_option("--local", dest="local", action="store_true", default=False)
    optionparser.add_option("--cluster", dest="cluster", action="store_true", default=False)
    # Specific options
    optionparser.add_option("-f", "--query_vcf_file", dest="query_vcf_file",
                            help="query_vcf_file")
    optionparser.add_option("-o", "--others_vcf_files", action="append", type="str",
                            dest="others_vcf_files",
                            help="others_vcf_files. Pass vcf files without .gz extention even if the file .gz is there")
    optionparser.add_option("-k", "--isec_with_knownSNPs", dest="isec_with_knownSNPs",
                            action="store_true", default=False,
                            help="others_vcf_files")
    optionparser.add_option("-c", "--do_complement", dest="do_complement",
                            action="store_true", default=False,
                            help="do complement or intersection. Defaults intersection")


    (options, args) = optionparser.parse_args()    

    config = ExomePipelineConfig()
    config.from_xml(options.config_file)
    analysis = ExomeAnalysisConfig()
    analysis.from_xml(options.analysis_file, config.output_dir)


    if not (options.local ^ options.cluster):
        print "Must set either --local or --cluster to run job"
    if options.local:
        jobrunfunc = run_local
    elif options.cluster:
        jobrunfunc = qsub_cac
    
    if options.isec_with_knownSNPs:
        knownSNPs_isec(analysis, config, [options.query_vcf_file], jobrunfunc, 
                       options.do_complement)
    else:
        if not (options.query_vcf_file and options.others_vcf_files):
            print "Must set both query_vcf_file and others_vcf_files"
        else:
            SNPs_isec(options.query_vcf_file, options.others_vcf_files, 
                      config, jobrunfunc, options.do_complement)
    
